import json
import ctypes
import os
import shutil
import random
import site
import sys
import threading
import time
import cv2
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import torch
import torchvision
import court 
import serial
# from tq import intersection,cross_point,horizontal_vertical,clear_duplicate_lines,find_Left_Right


CONF_THRESH = 0.25
IOU_THRESHOLD = 0.4

def gstreamer_pipeline(
    capture_width=800, #摄像头预捕获的图像宽度
    capture_height=600, #摄像头预捕获的图像高度
    display_width=800, #窗口显示的图像宽度
    display_height=600, #窗口显示的图像高度
    framerate=60,       #捕获帧率
    flip_method=0,      #是否旋转图像
):
    return (
        "nvarguscamerasrc ! "
       	"video/x-raw(memory:NVMM), "
        "width=(int)%d, height=(int)%d, "
        "format=(string)NV12, framerate=(fraction)%d/1 ! "
        "nvvidconv flip-method=%d ! "
        "video/x-raw, width=(int)%d, height=(int)%d, format=(string)BGRx ! "
        "videoconvert ! "
        "video/x-raw, format=(string)BGR ! appsink max-buffers=1 drop=True "
        % (
            capture_width,
            capture_height,
            framerate,
            flip_method,
            display_width,
            display_height,
        )
    )

def undistort(frame):
    fx = 551.50821032
    cx = 401.25023504
    fy = 736.12692753
    cy = 213.432079
    k1, k2, p1, p2, k3 = -0.44003543, 0.34060944, 0.00762225, -0.00330469, -0.20257493


    k = np.array([
        [fx, 0, cx],
        [0, fy, cy],
        [0, 0, 1]
    ])
    # 畸变系数
    d = np.array([
        k1, k2, p1, p2, k3
    ])
    h, w = frame.shape[:2]
    mapx, mapy = cv2.initUndistortRectifyMap(k, d, None, k, (w, h), 5)
    return cv2.remap(frame, mapx, mapy, cv2.INTER_LINEAR)

class Colors:
    # Ultralytics color palette https://ultralytics.com/
    def __init__(self):
        # hex = matplotlib.colors.TABLEAU_COLORS.values()
        hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
               '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        self.palette = [self.hex2rgb('#' + c) for c in hex]
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  # rgb order (PIL)
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))


colors = Colors()  # create instance for 'from utils.pl
def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    """
    description: Plots one bounding box on image img,
                 this function comes from YoLov5 project.
    param: 
        x:      a box likes [x1,y1,x2,y2]
        img:    a opencv image object
        color:  color to draw rectangle, such as (0,255,0)
        label:  str
        line_thickness: int
    return:
        no return

    """
    tl = (
        line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    )  # line/font thickness
    #color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    # label+=str(c1)
    # label+=str(c2)
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled

        
        
        
        cv2.putText(
            img,
            label,
            (c1[0], c1[1] - 2),
            0,
            tl / 3,
            [225, 255, 255],
            thickness=tf,
            lineType=cv2.LINE_AA,
        )
        
categories = [
        "person","bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
        "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
        "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
        "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
        "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
        "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
        "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
        "hair drier", "toothbrush"
        ]

class YoLov5TRT(object):
    """
    description: A YOLOv5 class that warps TensorRT ops, preprocess and postprocess ops.
    """

    def __init__(self, engine_file_path):
        # Create a Context on this device,
        self.ctx = cuda.Device(0).make_context()
        stream = cuda.Stream()
        TRT_LOGGER = trt.Logger(trt.Logger.INFO)
        runtime = trt.Runtime(TRT_LOGGER)

        # Deserialize the engine from file
        with open(engine_file_path, "rb") as f:
            engine = runtime.deserialize_cuda_engine(f.read())


        context = engine.create_execution_context()

        host_inputs = []
        cuda_inputs = []
        host_outputs = []
        cuda_outputs = []
        bindings = []

        for binding in engine:
            print('bingding:', binding, engine.get_binding_shape(binding))
            size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
            dtype = trt.nptype(engine.get_binding_dtype(binding))
            # Allocate host and device buffers
            host_mem = cuda.pagelocked_empty(size, dtype)
            cuda_mem = cuda.mem_alloc(host_mem.nbytes)
            # Append the device buffer to device bindings.
            bindings.append(int(cuda_mem))
            # Append to the appropriate list.
            if engine.binding_is_input(binding):
                self.input_w = engine.get_binding_shape(binding)[-1]
                self.input_h = engine.get_binding_shape(binding)[-2]
                host_inputs.append(host_mem)
                cuda_inputs.append(cuda_mem)
            else:
                host_outputs.append(host_mem)
                cuda_outputs.append(cuda_mem)

        # Store
        self.stream = stream
        self.context = context
        self.engine = engine
        self.host_inputs = host_inputs
        self.cuda_inputs = cuda_inputs
        self.host_outputs = host_outputs
        self.cuda_outputs = cuda_outputs
        self.bindings = bindings
        self.batch_size = engine.max_batch_size

    def infer(self,img):
        
        threading.Thread.__init__(self)
        # Make self the active context, pushing it on top of the context stack.
        self.ctx.push()
        # Restore
        stream = self.stream
        context = self.context
        engine = self.engine
        host_inputs = self.host_inputs
        cuda_inputs = self.cuda_inputs
        host_outputs = self.host_outputs
        cuda_outputs = self.cuda_outputs
        bindings = self.bindings
        # Do image preprocess
        batch_origin_h = []
        batch_origin_w = []
        input_image, image_raw, origin_h, origin_w = self.preprocess_image(img)
        batch_origin_h.append(origin_h)
        batch_origin_w.append(origin_w)
        # Copy input image to host buffer
        np.copyto(host_inputs[0], input_image.ravel())
        start = time.time()
        # Transfer input data  to the GPU.
        cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
        # Run inference.
        context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle)
        # Transfer predictions back from the GPU.
        cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
        # Synchronize the stream
        stream.synchronize()
        end=time.time()
        self.ctx.pop()
        trt_outputs = host_outputs[0]
        fake_result = {}
        fake_result["algorithm_data"] = {
            "is_alert": False,
            "target_count": 0,
            "target_info": []
        }

        fake_result["model_data"] = {
            "objects": []
        }

        # Do postprocess
        for i in range(self.batch_size):
            result_boxes, result_scores, result_classid = self.post_process(
                trt_outputs[i * 6001: (i + 1) * 6001], batch_origin_h[i], batch_origin_w[i]
            )
            for j in range(len(result_boxes)):
                box = result_boxes[j]
                conf=result_scores[j]
                cls=result_classid[j]
                xyxy_list = torch.tensor(box).view(1, 4).view(-1).tolist()
                conf_list = conf.tolist()
                label = categories[int(cls)]
                color=colors(int(cls),True)
                print(box)

                global person_x,person_y,left_line,Right_line,G_x,G_y
                if label=="person":
                    person_x= (box[0]+box[2])/2
                    person_y= box[3]
                    print("person","x:",person_x,"y:",person_y)
                    coordinate_str="  "+"("+str(int(G_x))+","+str(int(G_y))+")"
                    plot_one_box(box,img,color=color,label="{}:{:.2f}{}".format(label,conf,coordinate_str))
                   
                    
                
                #plot_one_box(box,batch_image_raw[i],color=color,label="{}:{:.2f}".format(label, conf))
                #plot_one_box(box,batch_image_raw[i],label="{}:{:.2f}".format(label, conf))
                fake_result['model_data']['objects'].append({
                "xmin": int(xyxy_list[0]),
                "ymin": int(xyxy_list[1]),
                "xmax": int(xyxy_list[2]),
                "ymax": int(xyxy_list[3]),
                "confidence": conf_list,
                "name": label
                })
                if label == 'person':
                    fake_result['algorithm_data']['target_info'].append({
                        "xmin": int(xyxy_list[0]),
                        "ymin": int(xyxy_list[1]),
                        "xmax": int(xyxy_list[2]),
                        "ymax": int(xyxy_list[3]),
                        "confidence": conf_list,
                        "name": label
                        })
        fake_result['algorithm_data']['is_alert'] = True if len(
            fake_result['algorithm_data']['target_info']) > 0 else False
        fake_result['algorithm_data']["target_count"] = len(fake_result['algorithm_data']['target_info'])
        
        return json.dumps(fake_result, indent=4),end-start
        
    def bbox_iou(self, box1, box2, x1y1x2y2=True):
        """
        description: compute the IoU of two bounding boxes
        param:
            box1: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h))
            box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h))            
            x1y1x2y2: select the coordinate format
        return:
            iou: computed iou
        """
        if not x1y1x2y2:
            # Transform from center and width to exact coordinates
            b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
            b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
            b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
            b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
        else:
            # Get the coordinates of bounding boxes
            b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
            b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]

        # Get the coordinates of the intersection rectangle
        inter_rect_x1 = np.maximum(b1_x1, b2_x1)
        inter_rect_y1 = np.maximum(b1_y1, b2_y1)
        inter_rect_x2 = np.minimum(b1_x2, b2_x2)
        inter_rect_y2 = np.minimum(b1_y2, b2_y2)
        # Intersection area
        inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) * \
                     np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None)
        # Union Area
        b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
        b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)

        iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

        return iou

    def destroy(self):
        # Remove any context from the top of the context stack, deactivating it.
        self.ctx.pop()

    def non_max_suppression(self, prediction, origin_h, origin_w, conf_thres=0.5, nms_thres=0.4):
        """
        description: Removes detections with lower object confidence score than 'conf_thres' and performs
        Non-Maximum Suppression to further filter detections.
        param:
            prediction: detections, (x1, y1, x2, y2, conf, cls_id)
            origin_h: original image height
            origin_w: original image width
            conf_thres: a confidence threshold to filter detections
            nms_thres: a iou threshold to filter detections
        return:
            boxes: output after nms with the shape (x1, y1, x2, y2, conf, cls_id)
        """
        # Get the boxes that score > CONF_THRESH
        boxes = prediction[prediction[:, 4] >= conf_thres]
        # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2]
        boxes[:, :4] = self.xywh2xyxy(origin_h, origin_w, boxes[:, :4])
        # clip the coordinates
        boxes[:, 0] = np.clip(boxes[:, 0], 0, origin_w -1)
        boxes[:, 2] = np.clip(boxes[:, 2], 0, origin_w -1)
        boxes[:, 1] = np.clip(boxes[:, 1], 0, origin_h -1)
        boxes[:, 3] = np.clip(boxes[:, 3], 0, origin_h -1)
        # Object confidence
        confs = boxes[:, 4]
        # Sort by the confs
        boxes = boxes[np.argsort(-confs)]
        # Perform non-maximum suppression
        keep_boxes = []
        while boxes.shape[0]:
            large_overlap = self.bbox_iou(np.expand_dims(boxes[0, :4], 0), boxes[:, :4]) > nms_thres
            label_match = boxes[0, -1] == boxes[:, -1]
            # Indices of boxes with lower confidence scores, large IOUs and matching labels
            invalid = large_overlap & label_match
            keep_boxes += [boxes[0]]
            boxes = boxes[~invalid]
        boxes = np.stack(keep_boxes, 0) if len(keep_boxes) else np.array([])
        return boxes

    def xywh2xyxy(self, origin_h, origin_w, x):
        """
        description:    Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
        param:
            origin_h:   height of original image
            origin_w:   width of original image
            x:          A boxes numpy, each row is a box [center_x, center_y, w, h]
        return:
            y:          A boxes numpy, each row is a box [x1, y1, x2, y2]
        """
        y = np.zeros_like(x)
        r_w = self.input_w / origin_w
        r_h = self.input_h / origin_h
        if r_h > r_w:
            y[:, 0] = x[:, 0] - x[:, 2] / 2
            y[:, 2] = x[:, 0] + x[:, 2] / 2
            y[:, 1] = x[:, 1] - x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
            y[:, 3] = x[:, 1] + x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
            y /= r_w
        else:
            y[:, 0] = x[:, 0] - x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
            y[:, 2] = x[:, 0] + x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
            y[:, 1] = x[:, 1] - x[:, 3] / 2
            y[:, 3] = x[:, 1] + x[:, 3] / 2
            y /= r_h

        return y

    def post_process(self, output, origin_h, origin_w):
        """
        description: postprocess the prediction
        param:
            output:     A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...] 
            origin_h:   height of original image
            origin_w:   width of original image
        return:
            result_boxes: finally boxes, a boxes numpy, each row is a box [x1, y1, x2, y2]
            result_scores: finally scores, a numpy, each element is the score correspoing to box
            result_classid: finally classid, a numpy, each element is the classid correspoing to box
        """
        # Get the num of boxes detected
        num = int(output[0])
        # Reshape to a two dimentional ndarray
        pred = np.reshape(output[1:], (-1, 6))[:num, :]
        boxes = self.non_max_suppression(pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD)
        result_boxes = boxes[:, :4] if len(boxes) else np.array([])
        result_scores = boxes[:, 4] if len(boxes) else np.array([])
        result_classid = boxes[:, 5] if len(boxes) else np.array([])
        return result_boxes, result_scores, result_classid
        
    def preprocess_image(self, raw_bgr_image):
        """
        description: Convert BGR image to RGB,
                     resize and pad it to target size, normalize to [0,1],
                     transform to NCHW format.
        param:
            input_image_path: str, image path
        return:
            image:  the processed image
            image_raw: the original image
            h: original height
            w: original width
        """
        image_raw = raw_bgr_image
        h, w, c = image_raw.shape
        image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB)
        # Calculate widht and height and paddings
        r_w = self.input_w / w
        r_h = self.input_h / h
        if r_h > r_w:
            tw = self.input_w
            th = int(r_w * h)
            tx1 = tx2 = 0
            ty1 = int((self.input_h - th) / 2)
            ty2 = self.input_h - th - ty1
        else:
            tw = int(r_h * w)
            th = self.input_h
            tx1 = int((self.input_w - tw) / 2)
            tx2 = self.input_w - tw - tx1
            ty1 = ty2 = 0
        # Resize the image with long side while maintaining ratio
        image = cv2.resize(image, (tw, th))
        # Pad the short side with (128,128,128)
        image = cv2.copyMakeBorder(
            image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, None, (128, 128, 128)
        )
        image = image.astype(np.float32)
        # Normalize to [0,1]
        image /= 255.0
        # HWC to CHW format:
        image = np.transpose(image, [2, 0, 1])
        # CHW to NCHW format
        image = np.expand_dims(image, axis=0)
        # Convert the image to row-major order, also known as "C order":
        image = np.ascontiguousarray(image)
        return image, image_raw, h, w
    
def Person_Coords(person_x,person_y):
    if(len(left_line)):
        x1,y1,x2,y2,flag=left_line
        left_x=(person_y-y1)*(x2-x1)/(y2-y1)+x1
    if(len(Right_line)):
        x1,y1,x2,y2,flag=Right_line
        Right_x=(person_y-y1)*(x2-x1)/(y2-y1)+x1
    
    #算出当前人所在的边界两点
    # print(left_x,Right_x)

    pixel_to_CM=610/abs(Right_x-left_x)  #一个像素点对应的CM 610/(左线和右线的x距离)
    # X相对坐标
    G_x=(person_x-left_x)*pixel_to_CM  #（人的x-左线的x）* 一个像素点对应的CM


    pixel_to_CM=(410/abs(bottom_line[1]-top_line[1])) #一个像素点对应的CM 410/(底线X-顶线X)
    #Y相对坐标        
    G_y= ( ((lines[0][0][1]+lines[0][0][3])) /2-person_y)*pixel_to_CM
        #像素点间距                           
    
    G_x=int(G_x)
    G_y=int(G_y)
    
    
    # 0A A5  88 X Y A5 0A
    
   
    # label+=str("  ",G_x,",")
    # label+=str(G_y)
    

if __name__ == "__main__":

    person_x=0
    person_y=0
    G_x,G_y=0,0
    # load custom plugin and engine
    PLUGIN_LIBRARY = "build/libmyplugins.so"
    engine_file_path = "build/yolov5s.engine"

    # uart1 = serial.Serial("/dev/ttyTHS1", 115200, timeout=5)

    if len(sys.argv) > 1:
        engine_file_path = sys.argv[1]
    if len(sys.argv) > 2:
        PLUGIN_LIBRARY = sys.argv[2]

    ctypes.CDLL(PLUGIN_LIBRARY)
    # a YoLov5TRT instance
    yolov5_wrapper = YoLov5TRT(engine_file_path)
#------------------------------------------------------------------------------------------------------------------------ 推理视频
    
    video_path=video_path=gstreamer_pipeline(flip_method=0) #调用摄像头 0
    cap=cv2.VideoCapture(video_path)
    #cap.set(3,800)
    #cap.set(4,600)
    # out_mp4 = cv2.VideoWriter(r"home/race_car_out.mp4",cv2.VideoWriter_fourcc(*'mp4v'),30,(800,600))

    time.sleep(2)
    ret ,img = cap.read()
    img = undistort(img)
    lines =[]
    image_copy = img.copy()
    

    lines=court.Get_Lines(image_copy)
    # if(len(lines)>4):
    h_lines=[]
    s_lines=[]
    left_line=[]
    Right_line=[]
    top_line=[]
    bottom_line=[]
    points=[]

    #判断是竖线还是横线
    court_exist,h_lines,s_lines=court.horizontal_vertical(lines)
    if(court_exist==True):
        # 从大到小排序线
        court.sort(h_lines)
        #print(lines)
        # 清除多余的线
        court.clear_duplicate_lines(h_lines)
        #print(lines)

        lines=h_lines+s_lines

        print(lines)
        # 寻找球场线
        left_line,Right_line,top_line,bottom_line=court.find_course(lines)
        left_x=0
        Right_x=0
    

    try:
        while True:
            # idx_frame+=1
            ret ,img = cap.read()
            img = undistort(img)
            fake_result,use_time = yolov5_wrapper.infer(img)
            
            # out_mp4.write(img)
                # print(fake_result)
            #print(1.0/use_time)
            print('court_line',left_line,Right_line,top_line,bottom_line)
            # print(lines)
            # for count,line in enumerate(lines):
            #     x1, y1, x2, y2,flag = line[0]
            
            #     if line[0][4]=='h':
            #         cv2.line(img, (0, y1), (800, y2), (0, 0, 255), thickness=2)
            #     if line[0][4]=='s':
            #         cv2.line(img, (x1, y1), (x2, y2), (0, 255, 0), thickness=2)
            img=  court.draw_course(img,left_line,Right_line,top_line,bottom_line)
            
                
        #     #------------------------------将人在图片中的位置转为球场坐标----------------------------------------------
        #     # person_x=613
        #     # person_y=341
            
            #------------------------------------------------------------------------------------------------
            if (court_exist==True):
                Person_Coords(person_x,person_y)
                print("G:",G_x,G_y,"\n")
                
            if G_x>=0 and G_y>=0:
                # uart1.write(bytearray([0x0A, 0xA5,0x88,G_x>>8,G_x&0x00ff,G_y>>8,G_y&0x00ff,0xA5,0X0A]))
                print("send ok")

            cv2.imshow("win",img)
            key_code=cv2.waitKey(1)&0xff
            if( key_code == ord("s")):
                ret ,img = cap.read()
                img = undistort(img)
                lines =[]
                image_copy = img.copy()
                    
                lines=court.Get_Lines(image_copy)
                # if(len(lines)>4):
                h_lines=[]
                s_lines=[]
                #判断是竖线还是横线
                h_lines,s_lines=court.horizontal_vertical(lines)
                # 从大到小排序线
                court.sort(h_lines)

                # 清除多余的线
                court.clear_duplicate_lines(h_lines)

                lines=h_lines+s_lines

                # left_line=[]
                # Right_line=[]
                # top_line=[]
                # bottom_line=[]
                # points=[]
                print("site ok")
                # 寻找球场线
                left_line,Right_line,top_line,bottom_line=court.find_course(lines)
            if(  key_code==27):
                break
                # out.write(img)
                # if idx_frame==20:
                    # time.sleep(0.5)
                    # quit()  
        
        # out.release()
    finally:
        cap.release()
        yolov5_wrapper.destroy()
        cv2.destroyAllWindows()
        exit()
    
#-------------------------------------------------------------------------------------------------------------------------------  推理单张图片
    """
    img_path="/project/tensorrtx/yolov5/samples/bus.jpg"
    parent, filename = os.path.split(img_path)
    output_dir=os.getcwd()+"/output/"
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.mkdir(output_dir)
    img=cv2.imread(img_path)
    try:
        fake_result,use_time = yolov5_wrapper.infer(img)
        cv2.imwrite(output_dir+filename,img)
        print(fake_result)
        print(use_time*1000)

    finally:
        yolov5_wrapper.destroy()
    """
#-------------------------------------------------------------------------------------------------------------------------------  推理一个文件夹下的所有图片
    """
    work_dir="samples/"
    files=os.listdir(work_dir)
    output_dir=os.getcwd()+"/output/"
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.mkdir(output_dir)
    try:
        for file in files:
            file_path=os.path.join(work_dir,file)
            parent, filename = os.path.split(file_path)        
            img=cv2.imread(file_path)
            fake_result,use_time = yolov5_wrapper.infer(img)
            cv2.imwrite(output_dir+filename,img)
            #print(fake_result)
            print(use_time*1000)
    finally:
        yolov5_wrapper.destroy()

   """
#-------------------------------------------------------------------------------------------------------------------------------------
